import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from config import PROCESSED_DATA_DIR, BATCH_SIZE


class TranslationDataset(Dataset):
    """ 中英翻译数据集类 """

    def __init__(self, data_path):
        """ 初始化数据集 """
        self.data = pd.read_json(data_path, orient='records', lines=True).to_dict(orient='records')

    def __len__(self):
        """ 数据集样本数量 """
        return len(self.data)

    def __getitem__(self, index):
        """ 获取指定索引的数据样本 """
        input_tensor = torch.tensor(self.data[index]['zh'], dtype=torch.long)
        target_tensor = torch.tensor(self.data[index]['en'], dtype=torch.long)
        return input_tensor, target_tensor


def get_dataloader(train=True):
    """ 获取数据加载器 """
    data_path = PROCESSED_DATA_DIR / ('indexed_train.jsonl' if train else 'indexed_test.jsonl')
    dataset = TranslationDataset(data_path)
    return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)


if __name__ == '__main__':
    train_dataloader = get_dataloader()
    print(len(train_dataloader))  # 183

    test_dataloader = get_dataloader(train=False)
    print(len(test_dataloader))  # 46

    for inputs, target in train_dataloader:
        print(inputs.shape)  # [batch_size, seq_len]
        print(target.shape)  # [batch_size, seq_len]
        break
